{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "# Recurrent neural network and dynamical system analysis\n",
    "\n",
    "__Content Creator:__ Guangyu Robert Yang\n",
    "\n",
    "__Content Modified:__ Kai Chen\n",
    "\n",
    "In this tutorial, we will use supervised learning to train a recurrent neural network on a simple perceptual decision making task, and analyze the trained network using dynamical system analysis."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# # If on Google Colab, uncomment to install neurogym to use cognitive tasks\n",
    "! git clone https://github.com/neurogym/neurogym.git\n",
    "%cd neurogym/\n",
    "! pip install -e ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Defining a cognitive task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "import neurogym as ngym\n",
    "\n",
    "# Environment\n",
    "task = 'PerceptualDecisionMaking-v0'\n",
    "kwargs = {'dt': 100}\n",
    "seq_len = 100\n",
    "\n",
    "# Make supervised dataset\n",
    "dataset = ngym.Dataset(task, env_kwargs=kwargs, batch_size=16,\n",
    "                       seq_len=seq_len)\n",
    "\n",
    "# A sample environment from dataset\n",
    "env = dataset.env\n",
    "# Visualize the environment with 2 sample trials\n",
    "_ = ngym.utils.plot_env(env, num_trials=2)\n",
    "\n",
    "# Network input and output size\n",
    "input_size = env.observation_space.shape[0]\n",
    "output_size = env.action_space.n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Define a vanilla continuous-time recurrent network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "Here we will define a continuous-time neural network but discretize it in time using the Euler method.\n",
    "\\begin{align}\n",
    "    \\tau \\frac{d\\mathbf{r}}{dt} = -\\mathbf{r}(t) + f(W_r \\mathbf{r}(t) + W_x \\mathbf{x}(t) + \\mathbf{b}_r).\n",
    "\\end{align}\n",
    "\n",
    "This continuous-time system can then be discretized using the Euler method with a time step of $\\Delta t$, \n",
    "\\begin{align}\n",
    "    \\mathbf{r}(t+\\Delta t) = \\mathbf{r}(t) + \\Delta \\mathbf{r} = \\mathbf{r}(t) + \\frac{\\Delta t}{\\tau}[-\\mathbf{r}(t) + f(W_r \\mathbf{r}(t) + W_x \\mathbf{x}(t) + \\mathbf{b}_r)].\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# Define networks\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import init\n",
    "from torch.nn import functional as F\n",
    "import math\n",
    "\n",
    "\n",
    "class CTRNN(nn.Module):\n",
    "    \"\"\"Continuous-time RNN.\n",
    "\n",
    "    Args:\n",
    "        input_size: Number of input neurons\n",
    "        hidden_size: Number of hidden neurons\n",
    "\n",
    "    Inputs:\n",
    "        input: (seq_len, batch, input_size), network input\n",
    "        hidden: (batch, hidden_size), initial hidden activity\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_size, hidden_size, dt=None, **kwargs):\n",
    "        super().__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.tau = 100\n",
    "        if dt is None:\n",
    "            alpha = 1\n",
    "        else:\n",
    "            alpha = dt / self.tau\n",
    "        self.alpha = alpha\n",
    "        self.oneminusalpha = 1 - alpha\n",
    "\n",
    "        self.input2h = nn.Linear(input_size, hidden_size)\n",
    "        self.h2h = nn.Linear(hidden_size, hidden_size)\n",
    "\n",
    "    def init_hidden(self, input_shape):\n",
    "        batch_size = input_shape[1]\n",
    "        return torch.zeros(batch_size, self.hidden_size)\n",
    "\n",
    "    def recurrence(self, input, hidden):\n",
    "        \"\"\"Recurrence helper.\"\"\"\n",
    "        pre_activation = self.input2h(input) + self.h2h(hidden)\n",
    "        h_new = torch.relu(hidden * self.oneminusalpha +\n",
    "                           pre_activation * self.alpha)\n",
    "        return h_new\n",
    "\n",
    "    def forward(self, input, hidden=None):\n",
    "        \"\"\"Propogate input through the network.\"\"\"\n",
    "        if hidden is None:\n",
    "            hidden = self.init_hidden(input.shape).to(input.device)\n",
    "\n",
    "        output = []\n",
    "        steps = range(input.size(0))\n",
    "        for i in steps:\n",
    "            hidden = self.recurrence(input[i], hidden)\n",
    "            output.append(hidden)\n",
    "\n",
    "        output = torch.stack(output, dim=0)\n",
    "        return output, hidden\n",
    "\n",
    "\n",
    "class RNNNet(nn.Module):\n",
    "    \"\"\"Recurrent network model.\n",
    "\n",
    "    Args:\n",
    "        input_size: int, input size\n",
    "        hidden_size: int, hidden size\n",
    "        output_size: int, output size\n",
    "        rnn: str, type of RNN, lstm, rnn, ctrnn, or eirnn\n",
    "    \"\"\"\n",
    "    def __init__(self, input_size, hidden_size, output_size, **kwargs):\n",
    "        super().__init__()\n",
    "\n",
    "        # Continuous time RNN\n",
    "        self.rnn = CTRNN(input_size, hidden_size, **kwargs)\n",
    "        self.fc = nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rnn_activity, _ = self.rnn(x)\n",
    "        out = self.fc(rnn_activity)\n",
    "        return out, rnn_activity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Train the recurrent network on the decision-making task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "# Instantiate the network and print information\n",
    "hidden_size = 64\n",
    "net = RNNNet(input_size=input_size, hidden_size=hidden_size,\n",
    "             output_size=output_size, dt=env.dt)\n",
    "print(net)\n",
    "\n",
    "# Use Adam optimizer\n",
    "optimizer = optim.Adam(net.parameters(), lr=0.01)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "running_loss = 0\n",
    "running_acc = 0\n",
    "for i in range(1000):\n",
    "    inputs, labels = dataset()\n",
    "    inputs = torch.from_numpy(inputs).type(torch.float)\n",
    "    labels = torch.from_numpy(labels.flatten()).type(torch.long)\n",
    "\n",
    "    # in your training loop:\n",
    "    optimizer.zero_grad()   # zero the gradient buffers\n",
    "    output, _ = net(inputs)\n",
    "    output = output.view(-1, output_size)\n",
    "    loss = criterion(output, labels)\n",
    "    loss.backward()\n",
    "    optimizer.step()    # Does the update\n",
    "\n",
    "    running_loss += loss.item()\n",
    "    if i % 100 == 99:\n",
    "        running_loss /= 100\n",
    "        print('Step {}, Loss {:0.4f}'.format(i+1, running_loss))\n",
    "        running_loss = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Visualize neural activity for in sample trials\n",
    "\n",
    "We will run the network for 100 sample trials, then visual the neural activity trajectories in a PCA space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "env.reset(no_step=True)\n",
    "perf = 0\n",
    "num_trial = 100\n",
    "activity_dict = {}\n",
    "trial_infos = {}\n",
    "for i in range(num_trial):\n",
    "    env.new_trial()\n",
    "    ob, gt = env.ob, env.gt\n",
    "    inputs = torch.from_numpy(ob[:, np.newaxis, :]).type(torch.float)\n",
    "    action_pred, rnn_activity = net(inputs)\n",
    "    rnn_activity = rnn_activity[:, 0, :].detach().numpy()\n",
    "    activity_dict[i] = rnn_activity\n",
    "    trial_infos[i] = env.trial\n",
    "\n",
    "# Concatenate activity for PCA\n",
    "activity = np.concatenate(list(activity_dict[i] for i in range(num_trial)), axis=0)\n",
    "print('Shape of the neural activity: (Time points, Neurons): ', activity.shape)\n",
    "\n",
    "# Print trial informations\n",
    "for i in range(5):\n",
    "    print('Trial ', i, trial_infos[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# Compute PCA and visualize\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "pca = PCA(n_components=2)\n",
    "pca.fit(activity)\n",
    "# print('Shape of the projected activity: (Time points, PCs): ', activity_pc.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "Transform individual trials and Visualize in PC space based on ground-truth color. We see that the neural activity is organized by stimulus ground-truth in PC1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, sharey=True, sharex=True, figsize=(6, 3))\n",
    "for i in range(num_trial):\n",
    "    activity_pc = pca.transform(activity_dict[i])\n",
    "    trial = trial_infos[i]\n",
    "    color = 'red' if trial['ground_truth'] == 0 else 'blue'\n",
    "    _ = ax1.plot(activity_pc[:, 0], activity_pc[:, 1], 'o-', color=color)\n",
    "\n",
    "    if i < 3:\n",
    "        _ = ax2.plot(activity_pc[:, 0], activity_pc[:, 1], 'o-', color=color)\n",
    "\n",
    "ax1.set_xlabel('PC 1')\n",
    "ax1.set_ylabel('PC 2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Dynamical system analysis\n",
    "\n",
    "### Search for approximate fixed points\n",
    "Here we search for approximate fixed points and visualize them in the same PC space. In a generic dynamical system,\n",
    "\\begin{align}\n",
    "    \\frac{d\\mathbf{x}}{dt} = F(\\mathbf{x}),\n",
    "\\end{align}\n",
    "We can search for fixed points by doing the optimization\n",
    "\\begin{align}\n",
    "    \\mathrm{argmin}_{\\mathbf{x}} |F(\\mathbf{x})|^2.\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# Freeze for parameters in the recurrent network\n",
    "for param in net.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "batch_size = 64\n",
    "\n",
    "# Inputs should be the 0-coherence mean input during stimulus period\n",
    "# This will be task-specific\n",
    "input = np.tile([1, 0.5, 0.5], (batch_size, 1))\n",
    "input = torch.tensor(input, dtype=torch.float32)\n",
    "\n",
    "# Here hidden activity is the variable to be optimized\n",
    "# Initialized randomly for search in parallel (activity all positive)\n",
    "hidden = torch.tensor(np.random.rand(batch_size, hidden_size)*3,\n",
    "                      requires_grad=True, dtype=torch.float32)\n",
    "\n",
    "# Use Adam optimizer\n",
    "optimizer = optim.Adam([hidden], lr=0.001)\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "running_loss = 0\n",
    "for i in range(10000):\n",
    "    optimizer.zero_grad()   # zero the gradient buffers\n",
    "\n",
    "    # Take the one-step recurrent function from the trained network\n",
    "    new_h = net.rnn.recurrence(input, hidden)\n",
    "    loss = criterion(new_h, hidden)\n",
    "    loss.backward()\n",
    "    optimizer.step()    # Does the update\n",
    "\n",
    "    running_loss += loss.item()\n",
    "    if i % 1000 == 999:\n",
    "        running_loss /= 1000\n",
    "        print('Step {}, Loss {:0.4f}'.format(i+1, running_loss))\n",
    "        running_loss = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Visualize the found approximate fixed points.\n",
    "\n",
    "We see that they found an approximate line attrator, corresponding to our PC1, along which evidence is integrated during the stimulus period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "fixedpoints = hidden.detach().numpy()\n",
    "print(fixedpoints.shape)\n",
    "\n",
    "# Plot in the same space as activity\n",
    "plt.figure()\n",
    "for i in range(5):\n",
    "    activity_pc = pca.transform(activity_dict[i])\n",
    "    trial = trial_infos[i]\n",
    "    color = 'red' if trial['ground_truth'] == 0 else 'blue'\n",
    "    plt.plot(activity_pc[:, 0], activity_pc[:, 1], 'o-',\n",
    "             color=color, alpha=0.1)\n",
    "\n",
    "# Fixed points are shown in cross\n",
    "fixedpoints_pc = pca.transform(fixedpoints)\n",
    "plt.plot(fixedpoints_pc[:, 0], fixedpoints_pc[:, 1], 'x')\n",
    "\n",
    "plt.xlabel('PC 1')\n",
    "plt.ylabel('PC 2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Computing the Jacobian and finding the line attractor\n",
    "\n",
    "First we will compute the Jacobian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# index of fixed point to focus on\n",
    "# choose one close to center by sorting PC1\n",
    "i_fp = np.argsort(fixedpoints[:, 0])[int(fixedpoints.shape[0]/2)]\n",
    "\n",
    "fp = torch.from_numpy(fixedpoints[i_fp])\n",
    "fp.requires_grad = True\n",
    "\n",
    "# Inputs should be the 0-coherence mean input during stimulus period\n",
    "# This will be task-specific\n",
    "input = torch.tensor([1, 0.5, 0.5], dtype=torch.float32)\n",
    "deltah = net.rnn.recurrence(input, fp) - fp\n",
    "\n",
    "# w1 = torch.randn((64,64), requires_grad = False)\n",
    "# new_h = w1@fp\n",
    "\n",
    "jacT = torch.zeros(hidden_size, hidden_size)\n",
    "for i in range(hidden_size):\n",
    "    output = torch.zeros(hidden_size)\n",
    "    output[i] = 1.\n",
    "    jacT[:,i] = torch.autograd.grad(deltah, fp, grad_outputs=output, retain_graph=True)[0]\n",
    "\n",
    "jac = jacT.detach().numpy().T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "Here we plot the direction of the eigenvector corresponding to the highest eigenvalue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "eigval, eigvec = np.linalg.eig(jac)\n",
    "vec = np.real(eigvec[:, np.argmax(eigval)])\n",
    "end_pts = np.array([+vec, -vec]) * 2\n",
    "end_pts = pca.transform(fp.detach().numpy() + end_pts)\n",
    "\n",
    "# Plot in the same space as activity\n",
    "plt.figure()\n",
    "for i in range(5):\n",
    "    activity_pc = pca.transform(activity_dict[i])\n",
    "    trial = trial_infos[i]\n",
    "    color = 'red' if trial['ground_truth'] == 0 else 'blue'\n",
    "    plt.plot(activity_pc[:, 0], activity_pc[:, 1], 'o-',\n",
    "             color=color, alpha=0.1)\n",
    "\n",
    "# Fixed points are shown in cross\n",
    "fixedpoints_pc = pca.transform(fixedpoints)\n",
    "plt.plot(fixedpoints_pc[:, 0], fixedpoints_pc[:, 1], 'x')\n",
    "\n",
    "# Line attractor\n",
    "plt.plot(end_pts[:, 0], end_pts[:, 1])\n",
    "\n",
    "plt.xlabel('PC 1')\n",
    "plt.ylabel('PC 2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# Plot distribution of eigenvalues in a 2-d real-imaginary plot\n",
    "plt.figure()\n",
    "plt.scatter(np.real(eigval), np.imag(eigval))\n",
    "plt.plot([0, 0], [-1, 1], '--')\n",
    "plt.xlabel('Real')\n",
    "plt.ylabel('Imaginary')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "# Supplementary Materials\n",
    "\n",
    "Code for making publication quality figures as it appears in the paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# Convert information into pandas dataframe\n",
    "import pandas as pd\n",
    "df = pd.DataFrame()\n",
    "for i in range(len(trial_infos)):\n",
    "    df = df.append(trial_infos[i], ignore_index=True)\n",
    "\n",
    "# Example selection of conditions\n",
    "print(df[(df['coh']==6.4) & (df['ground_truth']==1.0)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "plot_fp = False\n",
    "\n",
    "# Plot in the same space as activity\n",
    "fig = plt.figure(figsize=(3, 3))\n",
    "ax = fig.add_axes([0.2, 0.2, 0.6, 0.6])\n",
    "\n",
    "colors = np.array([[27,158,119], [117,112,179], [217,95,2]])/255.\n",
    "\n",
    "# Search for two trials with similar conditions\n",
    "cohs = np.unique(df['coh'])\n",
    "cohs = [c for c in cohs if c > 0]  # remove 0\n",
    "\n",
    "color_intensity = [0.4, 0.7, 1.0, 1.3]\n",
    "\n",
    "for ground_truth in [0, 1]:\n",
    "    if ground_truth == 0:\n",
    "        cohs_ = cohs[::-1]\n",
    "        color_intensity_ = color_intensity[::-1]\n",
    "    else:\n",
    "        cohs_ = cohs\n",
    "        color_intensity_ = color_intensity\n",
    "    for i_coh, coh in enumerate(cohs_):\n",
    "        if plot_fp and coh != cohs[2]:\n",
    "            continue\n",
    "        trials = df[(df['coh']==coh) & (df['ground_truth']==ground_truth)].index\n",
    "        activity = np.mean(np.array([activity_dict[i] for i in trials]), axis=0)\n",
    "\n",
    "        activity_pc = pca.transform(activity)\n",
    "        color = colors[0] if ground_truth == 0 else colors[1]\n",
    "        color = color * color_intensity_[i_coh]\n",
    "        signed_coh = coh * (2*ground_truth - 1)\n",
    "        label = '{:0.1f}'.format(signed_coh)\n",
    "        plt.plot(activity_pc[:, 0], activity_pc[:, 1], 'o-',\n",
    "                 color=color, ms=3, markeredgecolor='none',\n",
    "                 lw=1, label=label)\n",
    "\n",
    "\n",
    "if plot_fp:\n",
    "    # Fixed points are shown in cross\n",
    "    color = colors[2]\n",
    "    fixedpoints_pc = pca.transform(fixedpoints)\n",
    "    plt.plot(fixedpoints_pc[:, 0], fixedpoints_pc[:, 1], 'x', ms=3, color=color, alpha=0.3)\n",
    "\n",
    "    # Line attractor\n",
    "    plt.plot(fixedpoints_pc[i_fp, 0], fixedpoints_pc[i_fp, 1], 'x', ms=5, color=color, lw=1)\n",
    "    plt.plot(end_pts[:, 0], end_pts[:, 1], color=color)\n",
    "else:\n",
    "    ax.legend(title='Coh', loc='upper left', bbox_to_anchor=(1.0, 1.0), frameon=False)\n",
    "\n",
    "\n",
    "plt.xlabel('PC 1', fontsize=7)\n",
    "plt.ylabel('PC 2', fontsize=7)\n",
    "\n",
    "plt.xlim([-5, 5])\n",
    "plt.ylim([-1, 5])\n",
    "\n",
    "# Beautification\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['left'].set_position(('data', -5))\n",
    "ax.spines['bottom'].set_position(('data', -1.5))\n",
    "\n",
    "from pathlib import Path\n",
    "if plot_fp:\n",
    "    fname = Path('lineattractors_dm')\n",
    "else:\n",
    "    fname = Path('rnndynamics_dm')\n",
    "fig.savefig(fname.with_suffix('.pdf'), transparent=True)\n",
    "fig.savefig(fname.with_suffix('.png'), dpi=300)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "W5_Tutorial2",
   "toc_visible": true
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
